import os
import torch
import numpy as np
import pandas as pd
from torch.utils.data import Dataset
from load_data import loaddata


class TurbineDataset(Dataset):
    def __init__(self, df:pd.DataFrame, seq, device, base="speed", target="dspeed"):
        super(TurbineDataset, self).__init__()
        self.target = target
        self.device = device
        self.seq = seq
        self.df = df
        self.data = torch.tensor(self.df.to_numpy(dtype=np.float32), dtype=torch.float32).to(self.device)
        self.targetid = self.gettargetid(self.target)
        self.baseid = self.gettargetid(base)

    def __getitem__(self, index):
        return self.data[index:index+self.seq], self.data[index+self.seq-1][self.baseid], self.data[index+self.seq][self.targetid]
    
    def __len__(self):
        return len(self.data) - self.seq
    
    def gettargetid(self, target):
        cols = self.df.columns
        r = 0
        for col in cols:
            if target == col:
                return r
            r += 1
        return r


if __name__ == '__main__':
    df = loaddata()
    trainrate = 0.8
    testrate = 0.1
    seq = 12
    device = "cpu"
    traindf = df[:int(len(df) * trainrate)]
    testdf = df[-int(len(df) * testrate):]
    valdf = df[len(traindf):-len(testdf)]
    trainset = TurbineDataset(traindf, seq, device)
    valset = TurbineDataset(valdf, seq, device)
    testset = TurbineDataset(testdf, seq, device)
    print(len(trainset))
    print(len(valset))
    print(len(testset))